/*
Copyright 2025 PerfXLab (Beijing) Technologies Co., Ltd.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
#include "test_common.hpp"
#include <iostream>
#include <vector>
#include <complex>
#include <random>
#include <cmath>
#include <iomanip>
#include <limits>
#include <algorithm>

// 将我们的基础类型转换为NIST的枚举值
blas_base_type toNistBase(OpenSpB_base_type base) {
    return (base == OPENSPBLAS_ZERO_BASE) ? blas_zero_base : blas_one_base;
}

// 单精度浮点数比较函数
bool compareFloat(float a, float b, float eps = 1e-5) {
    float diff = std::abs(a - b);
    float maxVal = std::max(std::abs(a), std::abs(b));
    if (maxVal > 1.0f) {
        return (diff / maxVal) <= eps;
    } else {
        return diff <= eps;
    }
}

// 双精度浮点数比较函数
bool compareDouble(double a, double b, double eps = 1e-12) {
    double diff = std::abs(a - b);
    double maxVal = std::max(std::abs(a), std::abs(b));
    if (maxVal > 1.0) {
        return (diff / maxVal) <= eps;
    } else {
        return diff <= eps;
    }
}

// 单精度复数比较函数
bool compareComplex(const std::complex<float>& a, const std::complex<float>& b, float eps = 1e-5) {
    return compareFloat(a.real(), b.real(), eps) && 
           compareFloat(a.imag(), b.imag(), eps);
}

// 双精度复数比较函数
bool compareComplexDouble(const std::complex<double>& a, const std::complex<double>& b, double eps = 1e-12) {
    return compareDouble(a.real(), b.real(), eps) && 
           compareDouble(a.imag(), b.imag(), eps);
}

// 随机数生成器
template<typename T>
T randomValue() {
    static std::mt19937 gen(42); // 固定种子以便重现
    static std::uniform_real_distribution<float> dist_float(-10.0f, 10.0f);
    static std::uniform_real_distribution<double> dist_double(-10.0, 10.0);
    
    if constexpr (std::is_same_v<T, float>) {
        return dist_float(gen);
    } else if constexpr (std::is_same_v<T, double>) {
        return dist_double(gen);
    } else if constexpr (std::is_same_v<T, std::complex<float>>) {
        return {dist_float(gen), dist_float(gen)};
    } else if constexpr (std::is_same_v<T, std::complex<double>>) {
        return {dist_double(gen), dist_double(gen)};
    }
}

// 单精度测试函数
bool test_saxpy(OpenSpB_base_type index_base, float alpha, int incy) {
    // 准备测试数据
    int nz = 10;
    std::vector<int> indx(nz);
    int base_val = (index_base == OPENSPBLAS_ZERO_BASE) ? 0 : 1;
    
    // 生成随机索引 (保证在有效范围内)
    std::mt19937 gen(42);
    std::uniform_int_distribution<int> dist_index(base_val, base_val + 99);
    for (int i = 0; i < nz; ++i) {
        indx[i] = dist_index(gen);
    }
    
    // 计算所需y向量长度
    int max_index = *std::max_element(indx.begin(), indx.end());
    int y_len = (max_index - base_val) * incy + 1;
    
    // 分配并初始化向量
    std::vector<float> x(nz);
    std::vector<float> y_perf(y_len);
    std::vector<float> y_nist(y_len);
    
    for (int i = 0; i < nz; ++i) {
        x[i] = randomValue<float>();
    }
    
    for (int i = 0; i < y_len; ++i) {
        y_perf[i] = randomValue<float>();
        y_nist[i] = y_perf[i];
    }
    
    // 调用两个实现
    OpenSpB_saxpy(nz, alpha, x.data(), indx.data(), y_perf.data(), incy, index_base);
    BLAS_susaxpy(nz, alpha, x.data(), indx.data(), y_nist.data(), incy, toNistBase(index_base));
    
    // 验证结果
    for (int idx : indx) {
        int pos = (idx - base_val) * incy;
        if (pos < 0 || pos >= y_len) continue;
        
        if (!compareFloat(y_perf[pos], y_nist[pos])) {
            std::cerr << "saxpy验证失败: index_base=" << (index_base == OPENSPBLAS_ZERO_BASE ? "0" : "1")
                      << " alpha=" << alpha << " incy=" << incy
                      << " 位置 " << pos << ": " << y_perf[pos] << " vs " << y_nist[pos] << std::endl;
            return false;
        }
    }
    
    return true;
}

// 双精度测试函数
bool test_daxpy(OpenSpB_base_type index_base, double alpha, int incy) {
    // 准备测试数据
    int nz = 10;
    std::vector<int> indx(nz);
    int base_val = (index_base == OPENSPBLAS_ZERO_BASE) ? 0 : 1;
    
    // 生成随机索引 (保证在有效范围内)
    std::mt19937 gen(42);
    std::uniform_int_distribution<int> dist_index(base_val, base_val + 99);
    for (int i = 0; i < nz; ++i) {
        indx[i] = dist_index(gen);
    }
    
    // 计算所需y向量长度
    int max_index = *std::max_element(indx.begin(), indx.end());
    int y_len = (max_index - base_val) * incy + 1;
    
    // 分配并初始化向量
    std::vector<double> x(nz);
    std::vector<double> y_perf(y_len);
    std::vector<double> y_nist(y_len);
    
    for (int i = 0; i < nz; ++i) {
        x[i] = randomValue<double>();
    }
    
    for (int i = 0; i < y_len; ++i) {
        y_perf[i] = randomValue<double>();
        y_nist[i] = y_perf[i];
    }
    
    // 调用两个实现
    OpenSpB_daxpy(nz, alpha, x.data(), indx.data(), y_perf.data(), incy, index_base);
    BLAS_dusaxpy(nz, alpha, x.data(), indx.data(), y_nist.data(), incy, toNistBase(index_base));
    
    // 验证结果
    for (int idx : indx) {
        int pos = (idx - base_val) * incy;
        if (pos < 0 || pos >= y_len) continue;
        
        if (!compareDouble(y_perf[pos], y_nist[pos])) {
            std::cerr << "daxpy验证失败: index_base=" << (index_base == OPENSPBLAS_ZERO_BASE ? "0" : "1")
                      << " alpha=" << alpha << " incy=" << incy
                      << " 位置 " << pos << ": " << y_perf[pos] << " vs " << y_nist[pos] << std::endl;
            return false;
        }
    }
    
    return true;
}

// 单精度复数测试函数
bool test_caxpy(OpenSpB_base_type index_base, const std::complex<float>& alpha, int incy) {
    // 准备测试数据
    int nz = 10;
    std::vector<int> indx(nz);
    int base_val = (index_base == OPENSPBLAS_ZERO_BASE) ? 0 : 1;
    
    // 生成随机索引 (保证在有效范围内)
    std::mt19937 gen(42);
    std::uniform_int_distribution<int> dist_index(base_val, base_val + 99);
    for (int i = 0; i < nz; ++i) {
        indx[i] = dist_index(gen);
    }
    
    // 计算所需y向量长度
    int max_index = *std::max_element(indx.begin(), indx.end());
    int y_len = (max_index - base_val) * incy + 1;
    
    // 分配并初始化向量
    std::vector<std::complex<float>> x(nz);
    std::vector<std::complex<float>> y_perf(y_len);
    std::vector<std::complex<float>> y_nist(y_len);
    
    for (int i = 0; i < nz; ++i) {
        x[i] = randomValue<std::complex<float>>();
    }
    
    for (int i = 0; i < y_len; ++i) {
        y_perf[i] = randomValue<std::complex<float>>();
        y_nist[i] = y_perf[i];
    }
    
    // 调用两个实现
    OpenSpB_caxpy(nz, &alpha, reinterpret_cast<const void*>(x.data()), 
                    indx.data(), reinterpret_cast<void*>(y_perf.data()), 
                    incy, index_base);
                    
    BLAS_cusaxpy(nz, &alpha, reinterpret_cast<const void*>(x.data()), 
                indx.data(), reinterpret_cast<void*>(y_nist.data()), 
                incy, toNistBase(index_base));
    
    // 验证结果
    for (int idx : indx) {
        int pos = (idx - base_val) * incy;
        if (pos < 0 || pos >= y_len) continue;
        
        if (!compareComplex(y_perf[pos], y_nist[pos])) {
            std::cerr << "caxpy验证失败: index_base=" << (index_base == OPENSPBLAS_ZERO_BASE ? "0" : "1")
                      << " alpha=(" << alpha.real() << "," << alpha.imag() << ")"
                      << " incy=" << incy
                      << " 位置 " << pos << ": (" << y_perf[pos].real() << "," << y_perf[pos].imag() 
                      << ") vs (" << y_nist[pos].real() << "," << y_nist[pos].imag() << ")" << std::endl;
            return false;
        }
    }
    
    return true;
}

// 双精度复数测试函数
bool test_zaxpy(OpenSpB_base_type index_base, const std::complex<double>& alpha, int incy) {
    // 准备测试数据
    int nz = 10;
    std::vector<int> indx(nz);
    int base_val = (index_base == OPENSPBLAS_ZERO_BASE) ? 0 : 1;
    
    // 生成随机索引 (保证在有效范围内)
    std::mt19937 gen(42);
    std::uniform_int_distribution<int> dist_index(base_val, base_val + 99);
    for (int i = 0; i < nz; ++i) {
        indx[i] = dist_index(gen);
    }
    
    // 计算所需y向量长度
    int max_index = *std::max_element(indx.begin(), indx.end());
    int y_len = (max_index - base_val) * incy + 1;
    
    // 分配并初始化向量
    std::vector<std::complex<double>> x(nz);
    std::vector<std::complex<double>> y_perf(y_len);
    std::vector<std::complex<double>> y_nist(y_len);
    
    for (int i = 0; i < nz; ++i) {
        x[i] = randomValue<std::complex<double>>();
    }
    
    for (int i = 0; i < y_len; ++i) {
        y_perf[i] = randomValue<std::complex<double>>();
        y_nist[i] = y_perf[i];
    }
    
    // 调用两个实现
    OpenSpB_zaxpy(nz, &alpha, reinterpret_cast<const void*>(x.data()), 
                    indx.data(), reinterpret_cast<void*>(y_perf.data()), 
                    incy, index_base);
                    
    BLAS_zusaxpy(nz, &alpha, reinterpret_cast<const void*>(x.data()), 
                indx.data(), reinterpret_cast<void*>(y_nist.data()), 
                incy, toNistBase(index_base));
    
    // 验证结果
    for (int idx : indx) {
        int pos = (idx - base_val) * incy;
        if (pos < 0 || pos >= y_len) continue;
        
        if (!compareComplexDouble(y_perf[pos], y_nist[pos])) {
            std::cerr << "zaxpy验证失败: index_base=" << (index_base == OPENSPBLAS_ZERO_BASE ? "0" : "1")
                      << " alpha=(" << alpha.real() << "," << alpha.imag() << ")"
                      << " incy=" << incy
                      << " 位置 " << pos << ": (" << y_perf[pos].real() << "," << y_perf[pos].imag() 
                      << ") vs (" << y_nist[pos].real() << "," << y_nist[pos].imag() << ")" << std::endl;
            return false;
        }
    }
    
    return true;
}

TEST(SparseBLASlevel2, axpy) {
    // 测试参数组合
    std::vector<OpenSpB_base_type> bases = {OPENSPBLAS_ZERO_BASE, OPENSPBLAS_ONE_BASE};
    std::vector<float> s_alphas = {0.0f, 1.0f, -1.0f, 2.5f, -2.5f};
    std::vector<double> d_alphas = {0.0, 1.0, -1.0, 2.5, -2.5};
    std::vector<std::complex<float>> c_alphas = {
        {0.0f, 0.0f}, {1.0f, 0.0f}, {0.0f, 1.0f}, {1.0f, 1.0f}, {-1.5f, 2.0f}
    };
    std::vector<std::complex<double>> z_alphas = {
        {0.0, 0.0}, {1.0, 0.0}, {0.0, 1.0}, {1.0, 1.0}, {-1.5, 2.0}
    };
    std::vector<int> incy_values = {1, 2, 3};
    
    int passed = 0, total = 0;
    
    // 测试单精度
    for (auto base : bases) {
        for (auto alpha : s_alphas) {
            for (auto incy : incy_values) {
                total++;
                if (test_saxpy(base, alpha, incy)) {
                    passed++;
                } else {
                    std::cerr << "测试失败: saxpy base=" << base 
                              << " alpha=" << alpha << " incy=" << incy << std::endl;
                }
            }
        }
    }
    
    // 测试双精度
    for (auto base : bases) {
        for (auto alpha : d_alphas) {
            for (auto incy : incy_values) {
                total++;
                if (test_daxpy(base, alpha, incy)) {
                    passed++;
                } else {
                    std::cerr << "测试失败: daxpy base=" << base 
                              << " alpha=" << alpha << " incy=" << incy << std::endl;
                }
            }
        }
    }
    
    // 测试单精度复数
    for (auto base : bases) {
        for (auto& alpha : c_alphas) {
            for (auto incy : incy_values) {
                total++;
                if (test_caxpy(base, alpha, incy)) {
                    passed++;
                } else {
                    std::cerr << "测试失败: caxpy base=" << base 
                              << " alpha=(" << alpha.real() << "," << alpha.imag() << ")"
                              << " incy=" << incy << std::endl;
                }
            }
        }
    }
    
    // 测试双精度复数
    for (auto base : bases) {
        for (auto& alpha : z_alphas) {
            for (auto incy : incy_values) {
                total++;
                if (test_zaxpy(base, alpha, incy)) {
                    passed++;
                } else {
                    std::cerr << "测试失败: zaxpy base=" << base 
                              << " alpha=(" << alpha.real() << "," << alpha.imag() << ")"
                              << " incy=" << incy << std::endl;
                }
            }
        }
    }
    
    // 输出总结
    std::cout << "All tests passed!\n";
}